"""
Jacobi Polynomials and Gauss-Lobatto Quadrature

Implements Jacobi polynomial evaluation, Gauss-Lobatto quadrature nodes and weights,
and Legendre derivative matrices for spectral element methods.
Port of the Python SEMGenerator functionality to Julia.
"""

"""
    jacobi_gl_nodes(alpha::Real, beta::Real, N::Int) -> Vector{Float64}

Compute Gauss-Lobatto-Jacobi nodes for polynomial order N.
"""
function jacobi_gl_nodes(alpha::Real, beta::Real, N::Int)
    if N == 0
        return [-1.0]
    elseif N == 1
        return [-1.0, 1.0]
    else
        # Interior nodes from Gauss-Jacobi quadrature
        x_int, _ = jacobi_gq_nodes_weights(alpha + 1, beta + 1, N - 2)
        # Combine with boundary points
        return vcat([-1.0], x_int, [1.0])
    end
end

"""
    jacobi_gl_weights(alpha::Real, beta::Real, N::Int) -> Vector{Float64}

Compute Gauss-Lobatto-Jacobi weights for polynomial order N.
"""
function jacobi_gl_weights(alpha::Real, beta::Real, N::Int)
    x = jacobi_gl_nodes(alpha, beta, N)
    w = zeros(Float64, N + 1)
    
    # Compute weights using Jacobi polynomial evaluation
    for i in 1:(N+1)
        pN = jacobi_polynomial_value(x[i], alpha, beta, N)
        w[i] = (2 * N + 1) / (N * (N + 1)) / (pN^2)
    end
    
    return w
end

"""
    jacobi_gq_nodes_weights(alpha::Real, beta::Real, N::Int) -> (Vector{Float64}, Vector{Float64})

Compute Gauss-Jacobi quadrature nodes and weights.
"""
function jacobi_gq_nodes_weights(alpha::Real, beta::Real, N::Int)
    if N == 0
        x = [-(alpha - beta) / (alpha + beta + 2)]
        w = [2.0]
        return x, w
    end
    
    # Build tridiagonal matrix for eigenvalue problem
    h1 = 2 .* (0:N) .+ (alpha + beta)
    J = diagm(0 => (-1/2) * (alpha^2 - beta^2) ./ (h1 .+ 2) ./ h1)
    
    # Off-diagonal elements
    for i in 1:N
        J[i, i+1] = 2 / (h1[i] + 2) * sqrt((i) * (i + alpha + beta) * 
                                           (i + alpha) * (i + beta) / 
                                           (h1[i] + 1) / (h1[i] + 3))
        J[i+1, i] = J[i, i+1]
    end
    
    # Handle special case when alpha + beta is small
    if abs(alpha + beta) < 10 * eps()
        J[1, 1] = 0.0
    end
    
    # Solve eigenvalue problem
    eigenvals, eigenvecs = eigen(J)
    idx = sortperm(eigenvals)
    eigenvals = eigenvals[idx]
    eigenvecs = eigenvecs[:, idx]
    
    # Nodes are eigenvalues
    x = eigenvals
    
    # Weights from first component of eigenvectors
    w = (eigenvecs[1, :].^2) * (2^(alpha + beta + 1)) / (alpha + beta + 1) * 
        gamma(alpha + 1) * gamma(beta + 1) / gamma(alpha + beta + 1)
    
    return x, abs.(w)
end

"""
    jacobi_polynomial_value(x::Real, alpha::Real, beta::Real, N::Int) -> Float64

Evaluate Jacobi polynomial P_N^{(alpha,beta)}(x) at point x.
"""
function jacobi_polynomial_value(x::Real, alpha::Real, beta::Real, N::Int)
    if N == 0
        gamma0 = (2^(alpha + beta + 1)) / (alpha + beta + 1) * 
                 gamma(alpha + 1) * gamma(beta + 1) / gamma(alpha + beta + 1)
        return 1.0 / sqrt(gamma0)
    elseif N == 1
        gamma0 = (2^(alpha + beta + 1)) / (alpha + beta + 1) * 
                 gamma(alpha + 1) * gamma(beta + 1) / gamma(alpha + beta + 1)
        gamma1 = (alpha + 1) * (beta + 1) / (alpha + beta + 3) * gamma0
        return (((alpha + beta + 2) * x / 2 + (alpha - beta) / 2) / sqrt(gamma1))
    else
        # Use three-term recurrence relation
        PL = zeros(Float64, N + 1)
        
        # Initialize
        gamma0 = (2^(alpha + beta + 1)) / (alpha + beta + 1) * 
                 gamma(alpha + 1) * gamma(beta + 1) / gamma(alpha + beta + 1)
        PL[1] = 1.0 / sqrt(gamma0)
        
        gamma1 = (alpha + 1) * (beta + 1) / (alpha + beta + 3) * gamma0
        PL[2] = (((alpha + beta + 2) * x / 2 + (alpha - beta) / 2) / sqrt(gamma1))
        
        # Recurrence relation
        aold = 2 / (2 + alpha + beta) * sqrt((alpha + 1) * (beta + 1) / (alpha + beta + 3))
        
        for i in 2:N
            h1 = 2 * (i-1) + alpha + beta
            anew = 2 / (h1 + 2) * sqrt(i * (i + alpha + beta) * (i + alpha) * (i + beta) / 
                                      (h1 + 1) / (h1 + 3))
            bnew = -(alpha^2 - beta^2) / h1 / (h1 + 2)
            PL[i+1] = 1 / anew * (-aold * PL[i-1] + (x - bnew) * PL[i])
            aold = anew
        end
        
        return PL[N+1]
    end
end

"""
    legendre_derivative_matrix(N::Int) -> (Matrix{Float64}, Vector{Float64}, Vector{Float64})

Compute Legendre differentiation matrix D, nodes r, and weights w.
"""
function legendre_derivative_matrix(N::Int)
    # Get Gauss-Lobatto-Legendre nodes and weights (alpha=beta=0)
    r = jacobi_gl_nodes(0, 0, N)
    w = jacobi_gl_weights(0, 0, N)
    
    Np = N + 1
    
    # Compute differentiation matrix
    # Use the algorithm from Shen's "Spectral Methods" book
    
    # Compute omega (products of distances for each node)
    omega = ones(Float64, Np)
    for i in 1:Np
        for j in 1:Np
            if i != j
                omega[i] *= (r[i] - r[j])
            end
        end
    end
    
    # Compute differentiation matrix
    D = zeros(Float64, Np, Np)
    for i in 1:Np
        for j in 1:Np
            if i != j
                D[i, j] = omega[j] / (omega[i] * (r[i] - r[j]))
            end
        end
    end
    
    # Diagonal elements: negative sum of off-diagonal elements
    for i in 1:Np
        D[i, i] = -sum(D[i, :])
    end
    
    return D, r, w
end

"""
    chebyshev_derivative_matrix(N::Int) -> (Matrix{Float64}, Vector{Float64})

Compute Chebyshev differentiation matrix and nodes.
"""
function chebyshev_derivative_matrix(N::Int)
    if N == 0
        D = zeros(1, 1)
        x = [1.0]
        return D, x
    end
    
    # Chebyshev nodes
    x = cos.(π * (0:N) / N)
    
    # Scaling factors
    c = ones(Float64, N + 1)
    c[1] = 2.0
    c[end] = 2.0
    c .*= (-1).^(0:N)
    
    # Distance matrix
    X = repeat(x, 1, N+1)
    dX = X' - X
    
    # Differentiation matrix
    D = (c ./ c') ./ (dX + I)  # Add identity to avoid division by zero
    D -= diagm(0 => sum(D, dims=2)[:])  # Diagonal correction
    
    return D, x
end

"""
    finite_difference_matrix_1d(N::Int) -> (Matrix{Float64}, Matrix{Float64}, Vector{Float64})

Create 1D finite difference matrices for first and second derivatives.
"""
function finite_difference_matrix_1d(N::Int)
    # Grid points
    x = collect(range(-1, 1, length=N+1))
    dx = 2.0 / N
    
    # First derivative matrix (central differences with one-sided at boundaries)
    D1 = zeros(Float64, N+1, N+1)
    for i in 2:N
        D1[i, i-1] = -1 / (2*dx)
        D1[i, i+1] = 1 / (2*dx)
    end
    
    # Boundary conditions for first derivative
    D1[1, 1] = -1/dx
    D1[1, 2] = 1/dx
    D1[N+1, N] = -1/dx
    D1[N+1, N+1] = 1/dx
    
    # Second derivative matrix (central differences)
    D2 = zeros(Float64, N+1, N+1)
    for i in 2:N
        D2[i, i-1] = 1 / dx^2
        D2[i, i] = -2 / dx^2
        D2[i, i+1] = 1 / dx^2
    end
    
    # Boundary conditions for second derivative (can be modified as needed)
    D2[1, :] = D2[2, :]  # Copy from interior
    D2[N+1, :] = D2[N, :]
    
    return D1, D2, x
end

"""
    create_tensor_product_operators(D1_1d::Matrix{Float64}, dim::Int) -> NTuple

Create tensor product differentiation operators for multi-dimensional problems.
"""
function create_tensor_product_operators(D1_1d::Matrix{Float64}, dim::Int)
    N = size(D1_1d, 1) - 1
    I_1d = Matrix{Float64}(I, N+1, N+1)
    
    if dim == 2
        # 2D operators
        Dx = kron(I_1d, D1_1d)  # ∂/∂x
        Dy = kron(D1_1d, I_1d)  # ∂/∂y
        return (Dx, Dy)
    elseif dim == 3
        # 3D operators  
        Dx = kron(kron(I_1d, I_1d), D1_1d)  # ∂/∂x
        Dy = kron(kron(I_1d, D1_1d), I_1d)  # ∂/∂y
        Dz = kron(kron(D1_1d, I_1d), I_1d)  # ∂/∂z
        return (Dx, Dy, Dz)
    else
        error("Unsupported dimension: $dim")
    end
end

"""
    apply_derivative_operator(field::Array, D1::Matrix{Float64}, direction::Int, dim::Int) -> Array

Apply differentiation operator along specified direction.
"""
function apply_derivative_operator(field::Array, D1::Matrix{Float64}, direction::Int, dim::Int)
    if dim == 2
        if direction == 1  # x-direction
            return D1 * field
        else  # y-direction (direction == 2)
            return field * D1'
        end
    elseif dim == 3
        # For 3D, we need more careful tensor operations
        result = similar(field)
        
        if direction == 1  # x-direction
            for k in axes(field, 3), j in axes(field, 2)
                result[:, j, k] = D1 * field[:, j, k]
            end
        elseif direction == 2  # y-direction
            for k in axes(field, 3), i in axes(field, 1)
                result[i, :, k] = D1 * field[i, :, k]
            end
        else  # z-direction (direction == 3)
            for j in axes(field, 2), i in axes(field, 1)
                result[i, j, :] = D1 * field[i, j, :]
            end
        end
        
        return result
    else
        error("Unsupported dimension: $dim")
    end
end

"""
    compute_laplacian(field::Array, D2_ops::Tuple, dim::Int) -> Array

Compute Laplacian of a field using second derivative operators.
"""
function compute_laplacian(field::Array, D2_ops::Tuple, dim::Int)
    if dim == 2
        Dx2, Dy2 = D2_ops
        # ∇²u = ∂²u/∂x² + ∂²u/∂y²
        return apply_derivative_operator(field, Dx2, 1, dim) + 
               apply_derivative_operator(field, Dy2, 2, dim)
    elseif dim == 3
        Dx2, Dy2, Dz2 = D2_ops
        # ∇²u = ∂²u/∂x² + ∂²u/∂y² + ∂²u/∂z²
        return apply_derivative_operator(field, Dx2, 1, dim) + 
               apply_derivative_operator(field, Dy2, 2, dim) + 
               apply_derivative_operator(field, Dz2, 3, dim)
    else
        error("Unsupported dimension: $dim")
    end
end

"""
    spectral_interpolation_matrix(r_old::Vector{Float64}, r_new::Vector{Float64}) -> Matrix{Float64}

Create interpolation matrix from old nodes to new nodes using Lagrange polynomials.
"""
function spectral_interpolation_matrix(r_old::Vector{Float64}, r_new::Vector{Float64})
    N_old = length(r_old) - 1
    N_new = length(r_new) - 1
    
    # Interpolation matrix
    I_matrix = zeros(Float64, N_new + 1, N_old + 1)
    
    for i in 1:(N_new+1)
        for j in 1:(N_old+1)
            I_matrix[i, j] = lagrange_basis_function(j-1, r_old, r_new[i])
        end
    end
    
    return I_matrix
end

"""
    lagrange_basis_function(j::Int, nodes::Vector{Float64}, x::Float64) -> Float64

Evaluate j-th Lagrange basis function at point x.
"""
function lagrange_basis_function(j::Int, nodes::Vector{Float64}, x::Float64)
    N = length(nodes) - 1
    value = 1.0
    
    for k in 0:N
        if k != j
            value *= (x - nodes[k+1]) / (nodes[j+1] - nodes[k+1])
        end
    end
    
    return value
end